import base64
from typing import List

from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.history_aware_retriever import create_history_aware_retriever
from langchain.chains.retrieval import create_retrieval_chain
from langchain_core.documents import Document
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.retrievers import RetrieverOutputLike
from langchain_core.runnables import Runnable
from langchain_core.vectorstores import VectorStoreRetriever


def trunc_msgs(msgs: List[BaseMessage]):
    """
    截断消息列表，保留最新的30条消息。
    """
    return msgs[-10:] if len(msgs) > 10 else msgs


def combine(docs: List[Document]):
    return "\n\n".join([doc.metadata.get("answer") for doc in docs])


def contextualize_chain(llm: BaseChatModel, system_prompt: str, retriever: VectorStoreRetriever) -> Runnable:
    question_answer_chain = create_stuff_documents_chain(llm, create_qa_prompt(system_prompt))
    chain = create_retrieval_chain(create_context_retriever(llm, retriever), question_answer_chain)
    return chain


def create_qa_prompt(system_prompt: str) -> ChatPromptTemplate:
    return ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])


def create_context_retriever(llm: BaseChatModel, retriever: VectorStoreRetriever) -> RetrieverOutputLike:
    contextualize_q_system_prompt = """
    给定聊天历史和最新的用户问题（可能引用聊天历史中的上下文），
    制定一个独立的问题，该问题可以在没有聊天历史的情况下被理解。
    不要回答问题，只是根据需要重新制定问题，否则按原样返回。
    """
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ])
    # 用于感知会话历史的上下文，再根据上下文来检索相关文档，提高精确性
    return create_history_aware_retriever(
        llm, retriever, contextualize_q_prompt
    )



def encode_image_to_base64(image_path):
    """将图片编码为base64格式"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')


def create_multimodal_message(text, image_path=None):
    """创建包含文本和图片的消息"""
    content = [{"type": "text", "text": text}]
    if image_path:
        # 本地图片文件
        base64_image = encode_image_to_base64(image_path)
        content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:image/jpeg;base64,{base64_image}"
            }
        })
    return HumanMessage(content=content)


